# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
"""
Backbone modules.
"""
from collections import OrderedDict

import torch
import torch.nn.functional as F
import torchvision
from torch import nn
from torchvision.models._utils import IntermediateLayerGetter
from typing import Dict, List

from util.misc import NestedTensor, is_main_process

from .position_encoding import build_position_encoding


class FrozenBatchNorm2d(torch.nn.Module):#定义冻结的 BatchNorm2d 类，继承自 nn.Module，用于固定批归一化的统计量和仿射参数
    """
    BatchNorm2d where the batch statistics and the affine parameters are fixed.

    Copy-paste from torchvision.misc.ops with added eps before rqsrt,
    without which any other models than torchvision.models.resnet[18,34,50,101]
    produce nans.
    """
    #初始化函数：n：输入通道数。
    #注册缓冲区 weight、bias、running_mean 和 running_var，分别初始化为全 1、全 0、全 0 和全 1
    def __init__(self, n):
        super(FrozenBatchNorm2d, self).__init__()
        self.register_buffer("weight", torch.ones(n))
        self.register_buffer("bias", torch.zeros(n))
        self.register_buffer("running_mean", torch.zeros(n))
        self.register_buffer("running_var", torch.ones(n))

    #从状态字典加载参数时，删除 num_batches_tracked 键（因为冻结的 BatchNorm 不需要跟踪批次数）。
    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
                              missing_keys, unexpected_keys, error_msgs):
        num_batches_tracked_key = prefix + 'num_batches_tracked'
        if num_batches_tracked_key in state_dict:
            del state_dict[num_batches_tracked_key]
        #调用父类的 _load_from_state_dict 方法加载其余参数。
        super(FrozenBatchNorm2d, self)._load_from_state_dict(
            state_dict, prefix, local_metadata, strict,
            missing_keys, unexpected_keys, error_msgs)

    def forward(self, x):#前向传播：
        # move reshapes to the beginning
        # to make it fuser-friendly
        w = self.weight.reshape(1, -1, 1, 1)#将 weight、bias、running_var 和 running_mean 重塑为 [1, C, 1, 1]。
        b = self.bias.reshape(1, -1, 1, 1)
        rv = self.running_var.reshape(1, -1, 1, 1)
        rm = self.running_mean.reshape(1, -1, 1, 1)
        eps = 1e-5
        scale = w * (rv + eps).rsqrt()#计算缩放因子 scale 和偏置 bias。
        bias = b - rm * scale
        return x * scale + bias#对输入 x 进行缩放和偏置。

#定义 BackboneBase 类，继承自 nn.Module，用于实现主干网络的基础功能。
class BackboneBase(nn.Module):

    def __init__(self, backbone: nn.Module, train_backbone: bool, num_channels: int, return_interm_layers: bool):
        super().__init__()
        #根据 train_backbone 决定是否冻结部分层的参数：如果 train_backbone 为 False，冻结所有参数。
        # 如果 train_backbone 为 True，只训练 layer2、layer3 和 layer4 的参数。
        for name, parameter in backbone.named_parameters():
            if not train_backbone or 'layer2' not in name and 'layer3' not in name and 'layer4' not in name:
                parameter.requires_grad_(False)
        #根据 return_interm_layers 决定返回的层：如果为 True，返回 layer1 到 layer4 的输出。
        #如果为 False，只返回 layer4 的输出。使用 IntermediateLayerGetter 提取指定层的输出。
        if return_interm_layers:
            return_layers = {"layer1": "0", "layer2": "1", "layer3": "2", "layer4": "3"}
        else:
            return_layers = {'layer4': "0"}
        self.body = IntermediateLayerGetter(backbone, return_layers=return_layers)
        self.num_channels = num_channels#保存 num_channels 为类属性。

    def forward(self, tensor_list: NestedTensor):#前向传播：
        xs = self.body(tensor_list.tensors)#提取输入张量的特征 xs。
        out: Dict[str, NestedTensor] = {}
        for name, x in xs.items():
            m = tensor_list.mask
            assert m is not None
            mask = F.interpolate(m[None].float(), size=x.shape[-2:]).to(torch.bool)[0]
            out[name] = NestedTensor(x, mask)#对掩码 mask 进行插值，使其与特征图尺寸匹配。
        return out#返回包含特征和掩码的字典。

#定义 Backbone 类，继承自 BackboneBase，用于实现 ResNet 主干网络。
class Backbone(BackboneBase):
    """ResNet backbone with frozen BatchNorm."""
    def __init__(self, name: str,
                 train_backbone: bool,
                 return_interm_layers: bool,
                 dilation: bool):
        backbone = getattr(torchvision.models, name)(#根据 name 获取 ResNet 模型：
            replace_stride_with_dilation=[False, False, dilation],#replace_stride_with_dilation：控制是否用空洞卷积替换步长。
            pretrained=is_main_process(), norm_layer=FrozenBatchNorm2d)#如果是主进程，加载预训练权重。使用 FrozenBatchNorm2d 作为归一化层。
        #根据 ResNet 类型设置 num_channels：resnet18 和 resnet34 为 512。其他 ResNet 为 2048
        num_channels = 512 if name in ('resnet18', 'resnet34') else 2048
        #调用父类初始化。
        super().__init__(backbone, train_backbone, num_channels, return_interm_layers)


class Joiner(nn.Sequential):#定义 Joiner 类，继承自 nn.Sequential，用于将主干网络和位置编码结合。
    def __init__(self, backbone, position_embedding):
        super().__init__(backbone, position_embedding)

    def forward(self, tensor_list: NestedTensor):#前向传播：
        xs = self[0](tensor_list)#调用主干网络提取特征 xs
        out: List[NestedTensor] = []## 存储主干网络输出的特征图
        pos = []## 存储各层级特征图的位置编码
        for name, x in xs.items():
            out.append(x)# 保存当前层特征图
            # position encoding
            pos.append(self[1](x).to(x.tensors.dtype))#生成位置编码并转换数据类型

        return out, pos#返回特征图和位置编码的列表。


def build_backbone(args):#根据 args 构建位置编码 position_embedding。
    position_embedding = build_position_encoding(args)
    train_backbone = args.lr_backbone > 0#根据 lr_backbone 决定是否训练主干网络
    return_interm_layers = args.masks#根据 masks 决定是否返回中间层输出。
    backbone = Backbone(args.backbone, train_backbone, return_interm_layers, args.dilation)#构建主干网络 backbone。
    model = Joiner(backbone, position_embedding)#将主干网络和位置编码结合为 Joiner 模型。
    model.num_channels = backbone.num_channels#设置模型的 num_channels 属性并返回。
    return model
